# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import ujson as json
import datetime
import functools

from coin.base.book.types import (TradeSide,
                                  create_market_trade_event,
                                  create_market_ticker_event,
                                  create_market_index_event)
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.okex_futures.book_builder import BookBuilder
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.okex.kr_rest.currency import OkexCurrency
from coin.exchange.base.ws.trade_aggregator import TradeAggregator
from coin.exchange.okex_futures.ws.limit_price_checker.all_products \
    import LimitPriceChecker


def get_product_from_channel_name(channel, feed_type, as_of_timestamp):
  # channel 'ok_sub_futureusd_btc_trade_this_week'
  assert feed_type in ['depth', 'trade', 'ticker']
  assert feed_type in channel
  assert channel.startswith('ok_sub_')
  assert isinstance(as_of_timestamp, int)
  quote = channel[13:16]
  base = channel[17:20]
  currency_pair = "%s_%s" % (base, quote)
  contract_type = channel[channel.index(feed_type) + len(feed_type) + len('_'):]
  current_time = datetime.datetime.fromtimestamp(as_of_timestamp / 1.e9)
  return OkexFuturesProduct.FromCurrencyPairAndContractType(currency_pair,
                                                            contract_type,
                                                            current_time)


@functools.lru_cache(maxsize=128)
def get_product_from_channel_name_cached_impl(channel, feed_type, as_of_timestamp):
  return get_product_from_channel_name(channel, feed_type, as_of_timestamp)


def get_product_from_channel_name_cached(channel, feed_type, as_of_timestamp):
  # channel 'ok_sub_futureusd_btc_trade_this_week'
  # Unix timestamp does not take leap second into consideration.
  as_of_timestamp -= (60 * 10**9)  # Substract one minutes
  as_of_timestamp = (as_of_timestamp // (1800 * 10**9)) * (1800 * 10**9)
  return get_product_from_channel_name_cached_impl(channel, feed_type, as_of_timestamp)


@functools.lru_cache(maxsize=32)
def get_index_underlying_product_from_channel_name_cached(channel):
  assert channel.endswith('index')
  assert channel.startswith('ok_sub_futureusd')
  base_str = channel[17:-6]
  base = OkexCurrency.FromStrNativeCurrency(base_str)
  quote = OkexCurrency.FromStrCurrency('USD')
  return OkexProduct(base, quote)


class OkexFuturesHandler(WebsocketHandler):
  def __init__(self,
               *,
               event_dist=None,
               feed_checker_config=None,
               agg_view=True,
               check_limit_price=False):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._trade_processor = TradeAggregator(handler=self, agg_view=agg_view)
    if check_limit_price:
      self._limit_price_checker = LimitPriceChecker()
    else:
      self._limit_price_checker = None

  @property
  def limit_price_checker(self):
    return self._limit_price_checker

  def add_check_product(self, product):
    if self._limit_price_checker is not None:
      self._limit_price_checker.add_product(product)
    super().add_check_product(product)

  def on_first_book(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())

    channel = message['channel']
    product = get_product_from_channel_name_cached(channel, 'depth', record.timestamp)
    symbol = product.native_symbol
    if not self.is_subscribed('book', symbol):
      return
    data = message['data']

    current_datetime = datetime.datetime.fromtimestamp(record.timestamp / 1.e9)
    book = BookBuilder(symbol, current_datetime, book_builder_observer=self._book_builder_observer)
    bids = data.get('bids', [])
    asks = data.get('asks', [])
    book.update_fullbook(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self._books[symbol] = book

    self.ready = True
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    if self._limit_price_checker is not None:
      self._limit_price_checker.on_book(product=product, book=book, timestamp=record.timestamp)

    self.publish_event('book', book.symbol, book)

  def on_diff(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())

    channel = message['channel']
    product = get_product_from_channel_name_cached(channel, 'depth', record.timestamp)
    symbol = product.native_symbol
    if symbol not in self._books:
      return

    data = message['data']

    book = self._books[symbol]
    bids = data.get('bids', [])
    asks = data.get('asks', [])
    book.update_diff(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    if self._limit_price_checker is not None:
      self._limit_price_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())
    channel = message['channel']
    product = get_product_from_channel_name_cached(channel, 'trade', record.timestamp)
    symbol = product.native_symbol
    if not self.is_subscribed('trade', symbol):
      return
    data = message['data']

    self._trade_processor.on_begin_trade()

    last_idx = len(data) - 1
    for i, trade_msg in enumerate(data):
      price = float(trade_msg[1])
      # Okex Futures sends double-counted qty
      qty = float(trade_msg[2]) / 2
      assert trade_msg[4] == 'ask' or trade_msg[4] == 'bid'
      side = TradeSide.TRADE_BUY_SIDE if trade_msg[4] == 'bid' \
             else TradeSide.TRADE_SELL_SIDE
      trade_event = create_market_trade_event(record.timestamp,
                                              symbol,
                                              price,
                                              qty,
                                              side,
                                              has_more_trade=(i < last_idx),
                                              topic_id=record.topic_id)
      self._trade_processor.on_mid_trade(trade_event)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)

    self._trade_processor.on_end_trade()

  def subscribe_ticker(self, symbols, callback):
    self.subscribe_event('ticker', symbols, callback)

  def on_ticker(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())
    channel = message['channel']
    product = get_product_from_channel_name_cached(channel, 'ticker', record.timestamp)
    symbol = product.native_symbol
    data = message['data']

    ticker_event = create_market_ticker_event(timestamp=record.timestamp,
                                              product=product,
                                              low=float(data['low']),
                                              high=float(data['high']),
                                              vol_24h=float(data['vol']),
                                              open_interest=float(data['hold_amount']),
                                              limit_low=float(data['limitLow']),
                                              limit_high=float(data['limitHigh']))
    if self._limit_price_checker is not None:
      self._limit_price_checker.on_ticker(product=product,
                                          ticker_event=ticker_event,
                                          timestamp=record.timestamp)
    if not self.is_subscribed('ticker', symbol):
      return
    self.publish_event('ticker', symbol, ticker_event)

  def subscribe_index(self, symbols, callback):
    self.subscribe_event('index', symbols, callback)

  def on_index(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())
    channel = message['channel']
    underlying_product = get_index_underlying_product_from_channel_name_cached(channel)
    native_symbol = underlying_product.native_symbol
    if not self.is_subscribed('index', native_symbol):
      return
    data = message['data']

    index_event = create_market_index_event(timestamp=record.timestamp,
                                            product=underlying_product,
                                            index=data['futureIndex'])
    self.publish_event('index', native_symbol, index_event)
